import numpy as np
import pywt
import matplotlib.pyplot as plt
import matplotlib

# 生成一个模拟的2048点振动加速度数据（你可以用你的实际数据替代）
# 示例：简单的正弦波加噪声
fs = 10000  # 采样频率 (Hz)
N = 2048  # 数据点数
t = np.arange(N) / fs
f0 = 300  # 基频 (Hz)
data = 0.1* np.sin(2 * np.pi * f0 * t) + 0.5 * np.random.randn(N)
#目标频段
freq_target_limited = [100,600]

# 选择小波函数
FreqBand = 1.5
FreqCenter=1.0
wavelet = f'cmor{FreqBand}-{FreqCenter}'  # 复小波
#wavelet = 'cmor'
scales = np.arange(1, 128)  # 尺度范围



# 小波变换
coefficients, frequencies = pywt.cwt(data, scales, wavelet, sampling_period=1/fs)

# 计算频谱（绝对值的平方表示功率谱）
power_spectrum = np.abs(coefficients) ** 2

# 计算平均功率谱
avg_power_spectrum = np.mean(power_spectrum, axis=1)

# 过滤出关心的频率区间
freq_mask = (frequencies >= freq_target_limited[0]) & (frequencies <= freq_target_limited[1])
filtered_frequencies = frequencies[freq_mask]
filtered_power_spectrum = avg_power_spectrum[freq_mask]

# 找到最大功率谱的频率
max_freq_index = np.argmax(filtered_power_spectrum)
base_freq = filtered_frequencies[max_freq_index]

# 输出基频
print(f"提取的基频是：{base_freq:.2f} Hz")

# 设置字体路径
matplotlib.rcParams['font.sans-serif'] = ['SimHei']  # SimHei 是常用的中文字体
matplotlib.rcParams['axes.unicode_minus'] = False  # 处理负号问题

plt.subplot(311)  # 绘制时域波形
plt.plot(t, data)
plt.xlabel('Time (s)')
plt.ylabel('Amplitude')
plt.title('Time Domain Signal')
plt.grid()


plt.subplot(312)  # 绘制FFT
# 计算 FFT
fft_values = np.fft.fft(data)
fft_freqs = np.fft.fftfreq(N, 1/fs)

# 取绝对值并且取前半部分
fft_magnitude = np.abs(fft_values)/len(data)
half_n = N // 2
fft_freqs = fft_freqs[:half_n]
fft_magnitude = fft_magnitude[:half_n]

plt.plot(fft_freqs, fft_magnitude)
plt.title('FFT')
plt.xlabel('Frequency (Hz)')
plt.ylabel('Power')
plt.grid()


plt.subplot(313)  # 绘制功率谱
plt.plot(frequencies, avg_power_spectrum)
plt.title(f'Power Spectrum 特征频率 is {base_freq:.2f} in {freq_target_limited}')
plt.xlabel('Frequency (Hz)')
plt.ylabel('Power')
plt.grid()

plt.tight_layout()  # 调整子图间距
plt.show()
